import cv2
import pytesseract
import numpy as np
import pandas as pd
from PIL import Image

def safe_get(lst, index, default=None):
    """安全获取列表元素，避免索引越界"""
    try:
        return lst[index]
    except IndexError:
        return default

# 修改后的预处理部分
def preprocess_image(img):
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 自适应直方图均衡化
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)
    
    # 自适应阈值二值化（替代OTSU）
    binary = cv2.adaptiveThreshold(enhanced, 255, 
                                  cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                  cv2.THRESH_BINARY_INV, 11, 2)
    
    # 形态学操作增强线条
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    processed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=2)
    
    return processed

def detect_grid_lines(binary_img):
    # 水平线检测（加大核大小）
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (80, 1))
    horizontal = cv2.morphologyEx(binary_img, cv2.MORPH_OPEN, horizontal_kernel, iterations=3)
    
    # 垂直线检测（加大核大小）
    vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 80))
    vertical = cv2.morphologyEx(binary_img, cv2.MORPH_OPEN, vertical_kernel, iterations=3)
    
    # 合并线条
    grid_lines = cv2.addWeighted(horizontal, 0.5, vertical, 0.5, 0.0)
    
    # 增强交叉点
    intersections = cv2.bitwise_and(horizontal, vertical)
    
    return grid_lines, intersections, horizontal_kernel, vertical_kernel


def backup_grid_detection(binary_img):
    # 水平投影
    horizontal_proj = np.sum(binary_img, axis=1)
    row_positions = np.where(horizontal_proj > horizontal_proj.max()*0.1)[0]
    
    # 垂直投影
    vertical_proj = np.sum(binary_img, axis=0)
    col_positions = np.where(vertical_proj > vertical_proj.max()*0.1)[0]
    
    return row_positions, col_positions


def table_image_to_excel(image_path, output_excel):
    try:
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            raise FileNotFoundError(f"无法读取图像: {image_path}")
        
        # 图像预处理
        processed = preprocess_image(img)
        
        # 尝试检测网格线
        grid_lines, intersections, horizontal_kernel, vertical_kernel = detect_grid_lines(processed)
        
        if np.sum(intersections) == 0:
            print("警告：常规网格检测失败，尝试备用方法")
            row_positions, col_positions = backup_grid_detection(processed)
            
            if len(row_positions) < 2 or len(col_positions) < 2:
                raise ValueError("备用方法也无法检测表格结构")
            
            # 使用备用方法创建网格
            grid = []
            for i in range(len(row_positions)-1):
                row_cells = []
                for j in range(len(col_positions)-1):
                    cell = img[row_positions[i]:row_positions[i+1], 
                              col_positions[j]:col_positions[j+1]]
                    row_cells.append(cell)
                grid.append(row_cells)
        else:
            # 查找表格区域
            contours, _ = cv2.findContours(grid_lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if not contours:
                raise ValueError("未检测到表格区域")
            
            best_cnt = max(contours, key=cv2.contourArea)
            x, y, w, h = cv2.boundingRect(best_cnt)
            table_roi = img[y:y+h, x:x+w]

            # 检测表格内部结构
            roi_gray = cv2.cvtColor(table_roi, cv2.COLOR_BGR2GRAY)
            _, roi_thresh = cv2.threshold(roi_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
            roi_horizontal = cv2.morphologyEx(roi_thresh, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
            roi_vertical = cv2.morphologyEx(roi_thresh, cv2.MORPH_OPEN, vertical_kernel, iterations=2)
            intersections = cv2.bitwise_and(roi_vertical, roi_horizontal)

            # 获取网格点
            points = np.column_stack(np.where(intersections > 0))
            if len(points) == 0:
                raise ValueError("无法检测表格网格")

            # 组织行和列
            rows = {}
            for point in points[np.argsort(points[:, 0])]:
                row, col = point[0], point[1]
                rows.setdefault(row, set()).add(col)

            # 确定行位置
            row_positions = sorted(rows.keys())
            if not row_positions:
                raise ValueError("未检测到表格行")
                
            # 合并相近的行
            merged_rows = [row_positions[0]]
            for i in range(1, len(row_positions)):
                if row_positions[i] - merged_rows[-1] > 10:
                    merged_rows.append(row_positions[i])
            
            # 确保至少有两行
            if len(merged_rows) < 2:
                merged_rows = [0, roi_thresh.shape[0]-1]

            # 构建单元格网格
            grid = []
            for i in range(len(merged_rows)-1):
                row_start = merged_rows[i]
                row_end = merged_rows[i+1]
                
                # 获取当前行的列位置
                cols = sorted(rows.get(row_start, set()))
                if len(cols) < 2:
                    continue
                    
                # 合并相近的列
                merged_cols = [cols[0]]
                for j in range(1, len(cols)):
                    if cols[j] - merged_cols[-1] > 10:
                        merged_cols.append(cols[j])
                
                # 创建行单元格
                row_cells = []
                for j in range(len(merged_cols)-1):
                    col_start = merged_cols[j]
                    col_end = merged_cols[j+1]
                    cell_roi = table_roi[row_start:row_end, col_start:col_end]
                    row_cells.append(cell_roi)
                
                if row_cells:
                    grid.append(row_cells)

        # OCR识别
        table_data = []
        lang_config = '--oem 3 --psm 6 -l chi_sim+eng'
        
        for row in grid:
            row_texts = []
            for cell in row:
                cell_processed = preprocess_cell(cell)
                text = pytesseract.image_to_string(
                    Image.fromarray(cell_processed),
                    config=lang_config
                ).strip()
                row_texts.append(''.join(ch for ch in text if ch.isprintable()))
            table_data.append(row_texts)

        # 保存为Excel
        if not table_data:
            raise ValueError("未识别到有效数据")
            
        max_cols = max(len(row) for row in table_data)
        for row in table_data:
            row.extend([''] * (max_cols - len(row)))
        
        pd.DataFrame(table_data).to_excel(output_excel, index=False, header=False)
        print(f"成功转换: {output_excel}")

    except FileNotFoundError as e:
        print(f"文件错误: {str(e)}")
    except ValueError as e:
        print(f"值错误: {str(e)}")
    except Exception as e:
        print(f"处理失败: {str(e)}")
        # 保存中间处理结果用于调试
        cv2.imwrite("debug_processed.png", processed)
        cv2.imwrite("debug_grid.png", grid_lines*255)


def preprocess_cell(cell_img):
    """预处理单元格图像"""
    if len(cell_img) == 0 or len(cell_img[0]) == 0:
        return np.zeros((10, 10), dtype=np.uint8) + 255
    
    gray = cv2.cvtColor(cell_img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    if np.mean(thresh) < 127:
        thresh = cv2.bitwise_not(thresh)
    
    kernel = np.ones((1, 1), np.uint8)
    processed = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    processed = cv2.copyMakeBorder(processed, 15, 15, 15, 15, cv2.BORDER_CONSTANT, value=255)
    return processed


if __name__ == "__main__":
    table_image_to_excel("table_screenshot.png", "output_table.xlsx")
